import os

import pymysql
from dotenv import load_dotenv
from py2neo import Graph

load_dotenv()
MYSQL_CONFIG = {
    "host": "localhost",
    "port": 3306,  # MySQL默认端口，根据实际情况修改
    "user": 'root',
    "password":os.getenv('MYSQL_PASSWORD'),
    "database": os.getenv('DATABASE'),
    "charset": "utf8mb4",
}


def connect_mysql():
    try:
        conn = pymysql.connect(**MYSQL_CONFIG)
        print('--->数据库连接成功<---')
        return conn
    except Exception as e:
        raise Exception(f"数据库连接失败: {str(e)}")


def get_table_structure(conn):
    """定义游标，操作数据库"""
    tables = {}  # 用tables存储表的结构信息
    with conn.cursor() as cursor:
        '''with是为了方便管理，在结束后会关闭cursor游标'''
        cursor.execute("SHOW TABLES")
        table_names = [row[0] for row in cursor.fetchall()]
        '''['fault_info', 'fault_solution', 'fault_type']'''

        if not table_names:
            raise Exception("--->未获取到表名，请确认已经建表<---")

        '''提取字段信息'''
        for table_name in table_names:
            cursor.execute(f'show columns from {table_name}')
            columns = cursor.fetchall()
            '''
            (('emp_no', 'int(11)', 'NO', '', None, ''), ('dept_no', 'char(4)', 'NO', '', None, ''), ('from_date', 'date', 'YES', '', None, ''), ('to_date', 'date', 'YES', '', None, ''))
            (('dept_no', 'char(4)', 'NO', 'PRI', None, ''), ('dept_name', 'varchar(40)', 'NO', 'UNI', None, ''))
            (('emp_no', 'int(11)', 'NO', 'PRI', None, ''), ('dept_no', 'char(4)', 'NO', 'PRI', None, ''), ('from_date', 'date', 'NO', '', None, ''), ('to_date', 'date', 'NO', '', None, ''))
            (('emp_no', 'int(11)', 'NO', '', None, ''), ('from_date', 'date', 'YES', '', None, ''), ('to_date', 'date', 'YES', '', None, ''))
            (('emp_no', 'int(11)', 'NO', 'PRI', None, ''), ('dept_no', 'char(4)', 'NO', 'PRI', None, ''), ('from_date', 'date', 'NO', '', None, ''), ('to_date', 'date', 'NO', '', None, ''))
            (('emp_no', 'int(11)', 'NO', 'PRI', None, ''), ('birth_date', 'date', 'NO', '', None, ''), ('first_name', 'varchar(14)', 'NO', '', None, ''), ('last_name', 'varchar(16)', 'NO', '', None, ''), ('gender', "enum('M','F')", 'NO', '', None, ''), ('hire_date', 'date', 'NO', '', None, ''))
            (('emp_no', 'int(11)', 'NO', 'PRI', None, ''), ('salary', 'int(11)', 'NO', '', None, ''), ('from_date', 'date', 'NO', 'PRI', None, ''), ('to_date', 'date', 'NO', '', None, ''))
            (('emp_no', 'int(11)', 'NO', 'PRI', None, ''), ('title', 'varchar(50)', 'NO', 'PRI', None, ''), ('from_date', 'date', 'NO', 'PRI', None, ''), ('to_date', 'date', 'YES', '', None, ''))
             '''

            fields = []
            for col in columns:
                '''
                field_name用来获取每个字段的字段名
                key_type用来判断该字段是否为主键或者外键
                '''
                field_name = col[0]
                key_type = col[3]
                if key_type == 'PRI':
                    fields.append(f"{field_name} [主键]")
                elif key_type == 'MUL':
                    fields.append(f"{field_name} [外键]")
                else:
                    fields.append(field_name)
                    '''
                    ['emp_no', 'dept_no', 'from_date', 'to_date']
                    ['dept_no [主键]', 'dept_name']
                    ['emp_no [主键]', 'dept_no [主键]', 'from_date', 'to_date']
                    ['emp_no', 'from_date', 'to_date']
                    ['emp_no [主键]', 'dept_no [主键]', 'from_date', 'to_date']
                    ['emp_no [主键]', 'birth_date', 'first_name', 'last_name', 'gender', 'hire_date']
                    ['emp_no [主键]', 'salary', 'from_date [主键]', 'to_date']
                    ['emp_no [主键]', 'title [主键]', 'from_date [主键]', 'to_date']
                    '''
            tables[table_name] = {
                'name': table_name,
                'fields': fields,
            }
            '''
            {'name': 'current_dept_emp', 'fields': ['emp_no', 'dept_no', 'from_date', 'to_date']}
            {'name': 'departments', 'fields': ['dept_no [主键]', 'dept_name']}
            {'name': 'dept_emp', 'fields': ['emp_no [主键]', 'dept_no [主键]', 'from_date', 'to_date']}
            {'name': 'dept_emp_latest_date', 'fields': ['emp_no', 'from_date', 'to_date']}
            {'name': 'dept_manager', 'fields': ['emp_no [主键]', 'dept_no [主键]', 'from_date', 'to_date']}
            {'name': 'employees', 'fields': ['emp_no [主键]', 'birth_date', 'first_name', 'last_name', 'gender', 'hire_date']}
            {'name': 'salaries', 'fields': ['emp_no [主键]', 'salary', 'from_date [主键]', 'to_date']}
            {'name': 'titles', 'fields': ['emp_no [主键]', 'title [主键]', 'from_date [主键]', 'to_date']}
            '''
    print(f"--->成功提取到{len(tables)}张表的结构信息<---")
    return tables


TABLE_RELATIONSHIPS = [
    {
        "from_table": "fault_type",
        "to_table": "fault_info",
        "description": "故障类型对应多个故障记录（一对多关联，一个故障类型可包含多条故障记录）",
        "field_relation": "type_id → fault_info.type_id"
    },
    {
        "from_table": "fault_type",
        "to_table": "fault_solution",
        "description": "故障类型对应多个解决方案（一对多关联，一个故障类型可匹配多条解决方案）",
        "field_relation": "type_id → fault_solution.type_id"
    },
    {
        "from_table": "fault_info",
        "to_table": "fault_solution",
        "description": "故障记录对应多个解决方案（一对多关联，一个故障可对应多条解决方案）",
        "field_relation": "fault_id → fault_solution.fault_id"
    }
]

NEO4J_CONFIG = {
    'uri': 'bolt://localhost:7687',
    'user': 'neo4j',
    'password': os.getenv('PASSWORD_NEO')
}


def connect_neo4j():
    try:
        graph = Graph(
            uri='bolt://localhost:7687',
            auth=('neo4j', 'Shao123.')
        )
        print('--->连接Neo4j成功<---')
        return graph
    except Exception as e:
        raise Exception(f"Neo4j连接失败:{e}")


def create_neo4j_constraints(graph):
    graph.run("CREATE CONSTRAINT IF NOT EXISTS FOR (t:Table) REQUIRE t.name IS UNIQUE")
    print("--->唯一性约束创建成功<---")


def create_table_nodes(graph, tables):
    for table in tables.values():
        graph.run(
            """
            merge (t:Table {name:$name}) set t.label = $label, t.fields = $fields
            """,
            name=table["name"],
            label=table["name"],
            fields=table["fields"],
        )
    print(f"--->成功创建{len(tables)}个节点<---")


def create_table_relationships(graph):
    for rel in TABLE_RELATIONSHIPS:
        graph.run(
            '''
            match (from:Table {name:$from_table})
            match (to:Table {name:$to_table}) 
            merge (from)-[r:REFERENCES {
            description:$description,
            field_relation:$field_relation
            }]->(to)
            ''',
            from_table=rel["from_table"],
            to_table=rel["to_table"],
            field_relation=rel["field_relation"],
            description=rel["description"],
        )
    print(f"--->已在Neo4j中创建 {len(TABLE_RELATIONSHIPS)} 条表关系<---")


def main():
    conn = connect_mysql()
    graph = connect_neo4j()
    table_structure = get_table_structure(conn)
    create_neo4j_constraints(graph)
    create_table_nodes(graph, table_structure)
    create_table_relationships(graph)
    print("\n🎉 Text2SQL语义图谱构建完成！")
    print("   可通过 http://localhost:7474 访问Neo4j浏览器查看图谱")
    print("   在Neo4j浏览器中执行 MATCH (n) RETURN n 可查看所有节点和关系")

if __name__ == "__main__":
    main()
